package edu.washington.cs.workflow.analysis;

import java.util.LinkedHashSet;
import java.util.Set;

import com.ibm.wala.classLoader.IClass;
import com.ibm.wala.classLoader.IMethod;
import com.ibm.wala.core.tests.callGraph.CallGraphTestUtil;
import com.ibm.wala.ipa.callgraph.AnalysisCache;
import com.ibm.wala.ipa.callgraph.AnalysisOptions;
import com.ibm.wala.ipa.callgraph.AnalysisScope;
import com.ibm.wala.ipa.callgraph.CGNode;
import com.ibm.wala.ipa.callgraph.CallGraph;
import com.ibm.wala.ipa.callgraph.CallGraphBuilder;
import com.ibm.wala.ipa.callgraph.CallGraphStats;
import com.ibm.wala.ipa.callgraph.Entrypoint;
import com.ibm.wala.ipa.callgraph.impl.Util;
import com.ibm.wala.ipa.callgraph.propagation.PointerAnalysis;
import com.ibm.wala.ipa.cha.ClassHierarchy;
import com.ibm.wala.ipa.cha.IClassHierarchy;
import com.ibm.wala.util.config.AnalysisScopeReader;
import com.ibm.wala.util.graph.traverse.DFSFinishTimeIterator;
import com.ibm.wala.util.io.FileProvider;

import edu.washington.cs.workflow.util.Utils;
import edu.washington.cs.workflow.util.WALAUtils;

public class AnalysisBuilder {
	
	public enum CG {RTA, ZeroCFA, ZeroContainerCFA,
		VanillaZeroOneCFA, ZeroOneCFA, ZeroOneContainerCFA, OneCFA, TwoCFA, CFA, TempZeroCFA}
	
	public enum UIFramework {SWING}
	
	public final String classPath;
	public final String mainClass;
	public final String[] appPackages;
	
	private CG type = CG.ZeroCFA;
	private String byPassFile = null;
	private int cfaprecision = -1;
	private String exclusionFile = "Java60SwingInclusion.txt";
	
	private AnalysisScope scope = null;
	private ClassHierarchy cha = null;
	private Iterable<Entrypoint> entrypoints = null;
	private CallGraphBuilder builder = null;
	private AnalysisOptions options = null;
	private CallGraph cg = null;
	
	private UIFramework framework = null;
	
	public AnalysisBuilder(String classPath, String mainClass) {
		this(classPath, mainClass, null);
	}
	
	public AnalysisBuilder(String classPath, String mainClass, String[] appPackages) {
		this.classPath = classPath;
		this.mainClass = mainClass;
		this.appPackages = appPackages;
	}
	
	public void setUIFramework(UIFramework framework) {
		Utils.checkNotNull(framework);
		this.framework = framework;
	}
	
	public void buildAnalysis() {
		try {
		  System.out.println("Using exclusion file: " + this.exclusionFile);
		  System.out.println("CG type: " + this.type);
		  
	      this.scope = AnalysisScopeReader.makeJavaBinaryAnalysisScope(this.classPath, (new FileProvider())
	          .getFile(exclusionFile));
	      this.cha = ClassHierarchy.make(this.scope);
	      this.entrypoints = com.ibm.wala.ipa.callgraph.impl.Util.makeMainEntrypoints(scope, cha, mainClass);
	      if(this.framework == UIFramework.SWING) {
	    	  //add all event handling method as the entry points
	    	  this.entrypoints = SwingHelper.makeEntryPoints(scope, cha, mainClass, appPackages);
	      }
	      this.options = CallGraphTestUtil.makeAnalysisOptions(scope, entrypoints);
	      this.builder = chooseCallGraphBuilder(options, new AnalysisCache(), cha, scope);
	      //build the call graph
	      System.out.println("Building call graph...");
	      System.out.println("Number of entry points: " + Utils.countIterable(this.entrypoints));
	      this.cg = this.builder.makeCallGraph(options, null);
	      System.err.println(CallGraphStats.getStats(this.cg));
		} catch (Throwable e) {
			throw new Error(e);
		}
	}
	
	//only call this after building the CG
	//Never call this
	public void setExcludeStringBuilder(boolean exclude) {
		System.err.println("+++ Do you really want to call this method to exclude the string.");
		DFSFinishTimeIterator.DIRTY_HACK = exclude;
	}
	
	public void setCGType(CG type) {
		this.type = type;
	}
	
	public void setExclusionFile(String fileName) {
		this.exclusionFile = fileName;
	}
	
	public void setCFAPrecision(int length) {
		this.cfaprecision = length;
	}
	
	public ClassHierarchy getClassHierarchy() {
		return this.cha;
	}
	
	public CallGraph getCallGraph() {
		return this.cg;
	}
	
	public PointerAnalysis getPointerAnalysis() {
		return this.builder.getPointerAnalysis();
	}
	
	public CGNode findCGNodeBySignature(String signature) {
		for(CGNode node : this.getCallGraph()) {
			if(node.getMethod().getSignature().equals(signature)) {
				return node;
			}
		}
		return null;
	}
	
	public Set<String> getAllAppMethods() {
		Set<String> methodSigSet = new LinkedHashSet<String>();
		ClassHierarchy cha = this.getClassHierarchy();
		Utils.checkNotNull(cha);
		for(IClass c : cha) {
			String clazzName = WALAUtils.getJavaFullClassName(c);
			boolean isAppClass = Utils.startWith(clazzName, this.appPackages);
			if(!isAppClass) {
				continue;
			}
			for(IMethod m : c.getDeclaredMethods()) {
				String sig = m.getSignature();
				methodSigSet.add(sig);
			}
		}
		return methodSigSet;
	}
	
	public Set<String> getAllSwingHandlers() {
		Utils.checkTrue(this.framework.equals(UIFramework.SWING));
		Set<String> appHandlerSigSet = new LinkedHashSet<String>();
		ClassHierarchy cha = this.getClassHierarchy();
		Utils.checkNotNull(cha);
		for(IClass c : cha) {
			String clazzName = WALAUtils.getJavaFullClassName(c);
			boolean isAppClass = Utils.startWith(clazzName, this.appPackages);
			if(!isAppClass) {
				continue;
			}
			for(IMethod m : c.getDeclaredMethods()) {
				boolean isHandler = SwingHelper.isEventHandlingMethod(m, cha);
				if(isHandler) {
					appHandlerSigSet.add(m.getSignature());
				}
			}
		}
		return appHandlerSigSet;
	}
	
	private CallGraphBuilder chooseCallGraphBuilder(AnalysisOptions options, AnalysisCache cache,
	      IClassHierarchy cha, AnalysisScope scope) {
	    CallGraphBuilder builder = null;
	    if(this.type == CG.ZeroCFA) {
			System.out.println("Using 0-CFA call graph");
			builder = Util.makeZeroCFABuilder(options, cache, cha, scope);
		} else if (this.type == CG.ZeroOneCFA) {
			System.out.println("Using 0-1-CFA call graph");
			builder = Util.makeVanillaZeroOneCFABuilder(options, cache, cha, scope);
		} else if (this.type == CG.ZeroContainerCFA) {
			System.out.println("Using 0-container-CFA call graph");
			builder = Util.makeVanillaZeroOneContainerCFABuilder(options, cache, cha, scope);
		} else if (this.type == CG.RTA) {
			System.out.println("Using RTA call graph");
			builder = Util.makeRTABuilder(options, cache, cha, scope);
		} else if (this.type == CG.ZeroOneContainerCFA) {
			System.out.println("Using 0-1-container-CFA call graph");
			builder = Util.makeZeroOneContainerCFABuilder(options, cache, cha, scope);
		} else if (this.type == CG.OneCFA) {
			System.out.println("Using 1-CFA call graph");
			builder = WALAUtils.makeOneCFABuilder(options,  cache, cha, scope);
		} else if (this.type == CG.TwoCFA) {
			System.out.println("Using 2-CFA call graph");
			builder = WALAUtils.makeCFABuilder(2, options,  cache, cha, scope);
		} else if (this.type == CG.CFA) { 
			if(this.cfaprecision < 2) {
			throw new RuntimeException("Please set cfa precision first.");
			}
			System.out.println("Use CFA with precision: " + this.cfaprecision);
			builder = WALAUtils.makeCFABuilder(this.cfaprecision, options,  cache, cha, scope);
		} else if (this.type == CG.TempZeroCFA) {
			System.out.println("Use Temp-0-CFA with 0 context precision ");
			builder = WALAUtils.makeCFABuilder(0, options, cache, cha, scope);
		}else {
			throw new RuntimeException("The CG type: " + type + " is unknonw");
		}
		assert builder != null;
		
		//add the bypass file
		if(this.byPassFile != null) {
			System.err.println("Use bypass file: " + this.byPassFile);
			Util.addBypassLogic(options, scope, Utils.class.getClassLoader(), this.byPassFile, cha);
		}
		
		return builder;
	}
}